package org.pi.common.vision.io

import org.pi.common.matrix.FloatMatrix

import org.pi.common.matrix.Api._
import java.awt.image.BufferedImage
import java.io.File
import java.awt.color.ColorSpace
import java.awt.image.ComponentColorModel
import java.awt.Transparency
import java.awt.image.DataBuffer

object ImageIO {

  def makeGrayArrayFromImage(im: java.awt.Image): FloatMatrix = {
    var array: FloatMatrix = null
    val w = im.getWidth(null)
    val h = im.getHeight(null)
    if (w != -1 && h != -1) {
      array = FloatMatrix(h,w);
      val bImage = new BufferedImage(w,h,BufferedImage.TYPE_BYTE_GRAY);
      val g = bImage.createGraphics();
      g.drawImage(im,0,0,null);
      val r = bImage.getRaster();
      val data = new Array[Byte](w*h);
      var i: Int = 0;
      var x: Int = 0;
      var y: Int = 0;
      r.getDataElements(0,0,w,h,data);
      y=0;
      while (y < h) {
    	  x=0;
    	  while (x < w) {
    		  array(y,x) = (data(i).toInt & 0xff).toFloat / 255.0f
    		  i+=1
    		  x+=1
    	  }
    	  y+=1
      }
    }
    array;
  }

  def makeRgbArrayFromImage(im: java.awt.Image): FloatMatrix = {
    var array: FloatMatrix = null
    val w: Int = im.getWidth(null)
    val h: Int = im.getHeight(null)
    if (w != -1 && h != -1) {
      array = FloatMatrix(h,w,3);
      val bImage = new BufferedImage(w,h,BufferedImage.TYPE_3BYTE_BGR);
      val g = bImage.createGraphics();
      g.drawImage(im,0,0,null);
      val r = bImage.getRaster();
      val data = new Array[Byte](w*h*3);
      var i: Int = 0
      var x: Int = 0
      var y: Int = 0
      var c: Int = 0
      r.getDataElements(0,0,w,h,data);
      y=0;
      while (y < h) {
    	  x=0;
    	  while (x < w) {
    		  c=0;
    		  while (c < 3) {
    			  array(y,x,c) = (data(i).toInt & 0xff).toFloat / 255.0f
    			  i += 1
    			  c += 1
    		  }
    		  x+=1;
    	  }
    	  y+=1;
      }
    }
    array
  }
  
  def loadGrayArrayFromFile(name: String) : FloatMatrix = {
    try {
      val bImage = javax.imageio.ImageIO.read(new File(name));
      makeGrayArrayFromImage(bImage);
    }
    catch {
      case ex:Exception => ex.printStackTrace();
      null;
    }
  }

  def loadRgbArrayFromFile(name: String) : FloatMatrix = {
    try {
      val bImage = javax.imageio.ImageIO.read(new File(name));
      makeRgbArrayFromImage(bImage);
    }
    catch {
      case ex:Exception => ex.printStackTrace();
      null;
    }
  }

  def saveGrayArray(name:String, array: FloatMatrix, autoscale:Boolean, _minval: Float, _maxval: Float) {
    try {
      var bImage = createBufferedImageFromGrayArray(array, autoscale, _minval, _maxval)
      var format = name.substring(name.lastIndexOf(".")+1)
      javax.imageio.ImageIO.write(bImage,format,new File(name));
    }
    catch {
      case ex:Exception => ex.printStackTrace()
    }
  }

  def saveRgbArray(name:String, array: FloatMatrix, autoscale:Boolean, _minval: Float, _maxval: Float) {
    try {
      var bImage = createBufferedImageFromRgbArray(array, autoscale, _minval, _maxval)
      var format = name.substring(name.lastIndexOf(".")+1)
      javax.imageio.ImageIO.write(bImage,format,new File(name));
    }
    catch {
      case ex:Exception => ex.printStackTrace()
    }
  }
  
  def createBufferedImageFromGrayArray(array: FloatMatrix, autoscale: Boolean, _minval: Float, _maxval: Float) = {
    var value: Float = 0
    var minval = _minval
    var maxval=_maxval
    try {
      val cs_gray = ColorSpace.getInstance(ColorSpace.CS_GRAY);
      val cm_gray = new ComponentColorModel(cs_gray, false, true, Transparency.OPAQUE, DataBuffer.TYPE_BYTE);
      val h = size(array, 0)
      val w = size(array, 1)
      val gray = cm_gray.createCompatibleWritableRaster(w,h);
      val temp = new Array[Byte](w*h);
      if (autoscale) {
        minval = min(array);
        maxval = max(array);
        if (minval==maxval) {
          minval-=1;
          maxval+=1;
        }
      }
      var i: Int = 0
      var x: Int = 0
      var y: Int = 0
      while (y < h) {
    	  x=0;
    	  while (x < w) {
	        value = array(y,x)
	        if (value<minval) value=minval;
	        if (value>maxval) value=maxval;
	        temp(i) = ( (value-minval)/(maxval-minval)*255.0f ).toByte;
	        i += 1
	        x += 1
    	  }
    	  y += 1
      }
      gray.setDataElements(0,0,w,h,temp);
      new BufferedImage(cm_gray, gray, true, null);
    }
    catch {
      case ex: Exception => ex.printStackTrace()
      null
    }
  }

  def createBufferedImageFromRgbArray(array: FloatMatrix, autoscale: Boolean, _minval: Float, _maxval: Float) = {
    var value: Float = 0
    var minval = _minval
    var maxval=_maxval
    try {
      val cs_rgb = ColorSpace.getInstance(ColorSpace.CS_LINEAR_RGB);
      val cm_rgb = new ComponentColorModel(cs_rgb, false, true, Transparency.OPAQUE, DataBuffer.TYPE_BYTE);
      val h = size(array, 0)
      val w = size(array, 1)
      val rgb = cm_rgb.createCompatibleWritableRaster(w,h);
      val temp = new Array[Byte](3*w*h);
      if (autoscale) {
        minval = min(array);
        maxval = max(array);
        if (minval==maxval) {
          minval-=1;
          maxval+=1;
        }
      }
      var i: Int = 0
      var x: Int = 0
      var y: Int = 0
      var c: Int = 0
      y=0;
      while (y < h) {
    	  x=0;
    	  while (x < w) {
    		  c=0;
    		  while (c < 3) {
		        value = array(y,x,c)
		        if (value<minval) value=minval;
		        if (value>maxval) value=maxval;
		        temp(i) = ( (value-minval)/(maxval-minval)*255.0f ).toByte;
		        i += 1
		        c += 1
    		  }
    		  x += 1
    	  }
    	  y += 1
      }
      rgb.setDataElements(0,0,w,h,temp);
      new BufferedImage(cm_rgb, rgb, true, null);
    }
    catch {
      case ex: Exception => ex.printStackTrace()
      null
    }
  }
/*
  def main(arg: Array[String]) = {
    println("ImageTool Demo")
    var filename = "data/bsp01_farbe_klein.png"
    var grayimage = loadGrayArrayFromFile(filename)
    println("gray / size of "+filename+" is "+size(grayimage))
    var rgbimage = loadRgbArrayFromFile(filename)
    println("rgb / size of "+filename+" is "+size(rgbimage))
    saveGrayArray("test.png",grayimage, false, 0, 1)
    saveGrayArray("test.jpg",grayimage, false, 0, 1)
    saveRgbArray("testrgb.jpg",rgbimage, false, 0, 1)
//    saveGrayArray("testR.png",rgbimage., false, 0, 1)
  }
*/
}
